##!/usr/bin/env python3
'''
Basic IPv4 router (static routing) in Python.
'''

import sys
import os
import time
from switchyard.lib.userlib import *
from collections import deque
class ARP_request:
    def __init__(self, pkt, nextIP, nextDev, seq):
        self.pkts = [pkt]
        self.retry = 0
        self.time = -1
        self.nextIP = nextIP 
        self.nextDev = nextDev
        self.seq = seq

    def sent(self):
        self.retry += 1
        self.time = time.time()


class Router(object):
    def __init__(self, net):
        self.net = net
        self.arp_table = {}                 #IP-Eth mapping for other nodes
        self.router_interfaces = {}         #IP-Eth mapping for the router's interfaces
        self.forwarding_table = []          #forwarding table <ip_network, next_ip, eth_name>
        self.queue_size = 0                 #seq number for next queue item
        self.arp_queue = {}    #queue to keep track of arp requests sent out
        self.intf_name = {}

        # build forwarding table & router interface map
        for intf in net.interfaces():
            self.router_interfaces[intf.ipaddr] = intf.ethaddr
            log_info(intf.ipaddr)
            fwdEntry = [IPv4Network('{}/{}'.format(intf.ipaddr, intf.netmask), strict=False), None, intf.name]
            self.forwarding_table.append(fwdEntry)
            self.intf_name[intf.name] = intf

        # build forwarding table from file
        try:
            f = open("forwarding_table.txt")
            for line in f:
                entry = line.split()
                fwdEntry = [IPv4Network('{}/{}'.format(entry[0], entry[1]), strict=False), IPv4Address(entry[2]), entry[3]]
                self.forwarding_table.append(fwdEntry)
        except FileNotFoundError:
            log_debug("File not found")

    def prinf_arp_table(self):
        # TODO: print arp_table
        for key, value in self.arp_table.items():
            log_info('{key}: {value}'.format(key=key, value=value))
        log_info("\n")

    def add_to_queue(self, pkt, nextIP, nextDev):
        if nextIP in self.arp_queue:
            self.arp_queue[nextIP].pkts.append(pkt)
        else:
            self.queue_size += 1
            self.arp_queue[nextIP] = ARP_request(pkt, nextIP, nextDev, self.queue_size)

    def match_forwarding_table(self, ipaddr):
        maxLen = 0
        ansEntry = None
        for fwdEntry in self.forwarding_table:
            if ipaddr in fwdEntry[0]:
                if fwdEntry[0].prefixlen > maxLen:
                    maxLen = fwdEntry[0].prefixlen
                    ansEntry = fwdEntry
        return ansEntry

    def make_ARP_request(self, ipaddr, dev): #self.make_ARP_request(item.nextIP, item.nextDev)
        intf = self.intf_name[dev]
        arpReq = create_ip_arp_request(intf.ethaddr, intf.ipaddr, ipaddr)
        self.net.send_packet(intf.name, arpReq)

    def make_echo_reply(self, echoReq):
        fwdEntry = self.match_forwarding_table(echoReq[IPv4].src)
        nextIP = fwdEntry[1]
        if nextIP is None:
            nextIP = echoReq[IPv4].src  # TODO: use dst as next ip
        nextDev = fwdEntry[2]

        # TODO: creat ethernet header
        mac = Ethernet()
        mac.src = self.intf_name[nextDev].ethaddr
        # TODO: creat ip header 
        ip = IPv4()
        ip.src = echoReq[IPv4].dst
        ip.dst = echoReq[IPv4].src
        ip.ttl = 64
        # TODO: creat icmp header
        icmp = ICMP()
        icmp.icmptype = ICMPType.EchoReply
        icmp.icmpdata.identifier = echoReq[ICMP].icmpdata.identifier
        icmp.icmpdata.sequence = echoReq[ICMP].icmpdata.sequence
        icmp.icmpdata.data = echoReq[ICMP].icmpdata.data
        # TODO: send packet
        echoReply = mac + ip + icmp
        if nextIP in self.arp_table:
            self.sending_packet(echoReply, nextDev, self.arp_table[nextIP])
        else:
            self.add_to_queue(echoReply, nextIP, nextDev)

    def create_icmp_error_pkt(self, pkt, icmptype, icmpcode):
        fwdEntry = self.match_forwarding_table(pkt[IPv4].src)
        nextIP = fwdEntry[1]
        nextDev = fwdEntry[2]
        if nextIP is None:
            nextIP = pkt[IPv4].src      # TODO: use dst as next ip

        # TODO: creat ethernet header
        mac = Ethernet()
        mac.src = self.intf_name[nextDev].ethaddr
        # TODO: creat ip header
        ip = IPv4()
        ip.src = self.intf_name[nextDev].ipaddr
        ip.dst = pkt[IPv4].src
        ip.ttl = 64
        # TODO: creat icmp header
        icmp = ICMP()
        icmp.icmptype = icmptype
        icmp.icmpcode = icmpcode
        del pkt[pkt.get_header_index(Ethernet)]
        icmp.icmpdata.data = pkt.to_bytes()[:28]
        icmp.icmpdata.origdgramlen = len(pkt)
        # TODO: send packet   
        icmpErrMsg = mac + ip + icmp
        self.prinf_arp_table()
        if nextIP in self.arp_table:
            self.sending_packet(icmpErrMsg, nextDev, self.arp_table[nextIP])
        else:
            self.add_to_queue(icmpErrMsg, nextIP, nextDev)
    
 
    def handle_arp_requests(self):
        for key, item in sorted(list(self.arp_queue.items()), key=lambda x:x[1].seq): #sort by sequence
            if time.time() > item.time + 1:
                if item.retry >= 5:
                    for pkt in item.pkts:
                        self.create_icmp_error_pkt(pkt,ICMPType.DestinationUnreachable, ICMPTypeCodeMap[ICMPType.DestinationUnreachable].HostUnreachable)
                    del self.arp_queue[key]
                    self.handle_arp_requests()
                    return
                else:
                    self.make_ARP_request(item.nextIP, item.nextDev)
                    item.sent()

    def sending_packet(self, pkt, dev, nextMac):
        log_info("send packet from {} to {}".format(pkt[IPv4].src, pkt[IPv4].dst))
        pkt[Ethernet].src = self.intf_name[dev].ethaddr
        pkt[Ethernet].dst = nextMac
        self.net.send_packet(dev, pkt)

    def router_main(self):    
        while True: 
            gotpkt = True
            try:
                timestamp, dev, pkt = self.net.recv_packet(timeout=1.0)
            except NoPackets:
                log_debug("No packets available in recv_packet")
                gotpkt = False
            except Shutdown:
                log_debug("Got shutdown signal")
                break

            if gotpkt:
                if (pkt.has_header(Arp)):
                    arp = pkt.get_header(Arp)
                     # TODO: get ARP request
                    if arp.operation == 1:
                        # TODO: add request information to arp table
                        if arp.senderprotoaddr not in self.arp_table.keys():
                            self.arp_table[arp.senderprotoaddr] = arp.senderhwaddr
                        if arp.targetprotoaddr in self.router_interfaces:
                            targetHW = self.router_interfaces[arp.targetprotoaddr]
                            arpReply = create_ip_arp_reply( 
                                    targetHW,
                                    arp.senderhwaddr, 
                                    arp.targetprotoaddr,
                                    arp.senderprotoaddr)
                            self.net.send_packet(dev, arpReply)
                    # TODO: get ARP reply
                    else:
                        self.arp_table[arp.senderprotoaddr] = arp.senderhwaddr
                        if arp.senderprotoaddr in self.arp_queue:
                            queueItem = self.arp_queue[arp.senderprotoaddr]
                            for pkt in queueItem.pkts:
                                self.sending_packet(pkt, queueItem.nextDev, self.arp_table[queueItem.nextIP])
                            del self.arp_queue[arp.senderprotoaddr]

                # TODO: get IP packet
                elif(pkt.has_header(IPv4)):
                    dstIP = pkt[IPv4].dst
                    fwdEntry = self.match_forwarding_table(dstIP)
                    pkt[IPv4].ttl -= 1
                    if fwdEntry is not None: 
                        if pkt[IPv4].ttl <= 0: # TODO: time exceed
                            self.create_icmp_error_pkt(pkt, ICMPType.TimeExceeded, ICMPTypeCodeMap[ICMPType.TimeExceeded].TTLExpired)
                        else:
                            nextIP = fwdEntry[1]
                            nextDev = fwdEntry[2]

                            if dstIP in self.router_interfaces:
                                icmp = pkt.get_header(ICMP)
                                if icmp is not None and icmp.icmptype == ICMPType.EchoRequest:
                                    self.make_echo_reply(pkt)
                                else:
                                    self.create_icmp_error_pkt(pkt, ICMPType.DestinationUnreachable, ICMPTypeCodeMap[ICMPType.DestinationUnreachable].PortUnreachable)
                                   
                            else:
                                if nextIP is None:
                                    nextIP = pkt[IPv4].dst    
                                if nextIP in self.arp_table:
                                    self.sending_packet(pkt, nextDev, self.arp_table[nextIP])
                                else:
                                    self.add_to_queue(pkt, nextIP, nextDev)
                    else:
                        self.prinf_arp_table()
                        self.create_icmp_error_pkt(pkt, ICMPType.DestinationUnreachable, ICMPTypeCodeMap[ICMPType.DestinationUnreachable].NetworkUnreachable)

            self.handle_arp_requests() 




def main(net):
    '''
    Main entry point for router.  Just create Router
    object and get it going.
    '''
    r = Router(net)
    r.router_main()
    net.shutdown()
